package middleware

import (
	"crypto/rand"
	"encoding/base64"
	"errors"
	"net/http"
	"regexp"
	"time"

	"github.com/gin-gonic/gin"
	"go.uber.org/zap"
)

const (
	// CSRFTokenLength CSRF令牌长度
	CSRFTokenLength = 32
	// CSRFTokenExpiration CSRF令牌过期时间
	CSRFTokenExpiration = 3600 * time.Second
	// CSRFHeaderName CSRF请求头名称
	CSRFHeaderName = "X-CSRF-Token"
	// CSRFCookieName CSRF Cookie名称
	CSRFCookieName = "csrf_token"
)

// CSRFConfig CSRF配置
type CSRFConfig struct {
	Enable         bool     // 是否启用CSRF保护
	ExcludeURLs    []string // 排除的URL列表
	HeaderName     string   // CSRF令牌请求头名称
	CookieName     string   // CSRF令牌Cookie名称
	CookieDomain   string   // Cookie域名
	CookiePath     string   // Cookie路径
	CookieMaxAge   int      // Cookie最大生命周期
	CookieHTTPOnly bool     // 是否设置HttpOnly标志
	CookieSecure   bool     // 是否设置Secure标志
}

// DefaultCSRFConfig 默认CSRF配置
func DefaultCSRFConfig() CSRFConfig {
	return CSRFConfig{
		Enable:         true,
		ExcludeURLs:    []string{"/api/v1/auth/login", "/api/v1/auth/register"},
		HeaderName:     CSRFHeaderName,
		CookieName:     CSRFCookieName,
		CookiePath:     "/",
		CookieMaxAge:   86400,
		CookieHTTPOnly: true,
		CookieSecure:   false,
	}
}

// CSRF CSRF中间件
func CSRF() gin.HandlerFunc {
	return CSRFWithConfig(DefaultCSRFConfig())
}

// 检查是否为排除的URL
func isExcludedURL(path string, excludeRegexps []*regexp.Regexp) bool {
	for _, r := range excludeRegexps {
		if r.MatchString(path) {
			return true
		}
	}
	return false
}

// CSRFWithConfig 使用自定义配置的CSRF中间件
func CSRFWithConfig(config CSRFConfig) gin.HandlerFunc {
	// 如果未启用CSRF保护，则返回空中间件
	if !config.Enable {
		return func(c *gin.Context) {
			c.Next()
		}
	}

	// 编译排除URL的正则表达式
	excludeRegexps := make([]*regexp.Regexp, 0, len(config.ExcludeURLs))
	for _, pattern := range config.ExcludeURLs {
		r, err := regexp.Compile(pattern)
		if err != nil {
			zap.L().Error("编译CSRF排除URL正则表达式失败", zap.String("pattern", pattern), zap.Error(err))
			continue
		}
		excludeRegexps = append(excludeRegexps, r)
	}

	return func(c *gin.Context) {
		// 检查是否为排除的URL
		path := c.Request.URL.Path
		if isExcludedURL(path, excludeRegexps) {
			zap.L().Debug("CSRF保护：URL被排除", zap.String("path", path))
			c.Next()
			return
		}

		// 对于安全方法（GET、HEAD、OPTIONS、TRACE），不需要验证CSRF令牌
		if isSecureMethod(c.Request.Method) {
			// 为GET请求生成CSRF令牌
			if c.Request.Method == http.MethodGet {
				// 获取或生成CSRF令牌
				token, err := GenerateCSRFToken(c)
				if err != nil {
					zap.L().Error("生成CSRF令牌失败", zap.Error(err))
					ResponseError(c, http.StatusInternalServerError, "生成CSRF令牌失败")
					c.Abort()
					return
				}

				// 设置CSRF令牌Cookie
				SetCSRFCookie(c, token, config)
			}
			c.Next()
			return
		}

		// 对于不安全方法（POST、PUT、DELETE、PATCH），需要验证CSRF令牌
		token, err := getCSRFToken(c, config.HeaderName)
		if err != nil {
			zap.L().Error("获取CSRF令牌失败", zap.Error(err), zap.String("path", path), zap.String("method", c.Request.Method))
			ResponseError(c, http.StatusForbidden, "CSRF验证失败：缺少令牌")
			c.Abort()
			return
		}

		// 验证CSRF令牌
		valid, err := validateCSRFToken(c, token)
		if err != nil || !valid {
			zap.L().Error("验证CSRF令牌失败", zap.Error(err), zap.String("path", path), zap.String("method", c.Request.Method))
			ResponseError(c, http.StatusForbidden, "CSRF验证失败：无效的令牌")
			c.Abort()
			return
		}

		c.Next()
	}
}

// isSecureMethod 判断是否为安全的HTTP方法
func isSecureMethod(method string) bool {
	return method == http.MethodGet ||
		method == http.MethodHead ||
		method == http.MethodOptions ||
		method == http.MethodTrace
}

// 生成CSRF令牌
func GenerateCSRFToken(c *gin.Context) (string, error) {
	// 生成随机令牌
	b := make([]byte, 32)
	_, err := rand.Read(b)
	if err != nil {
		return "", err
	}

	token := base64.StdEncoding.EncodeToString(b)
	zap.L().Debug("生成CSRF令牌", zap.String("token", token))

	// 将令牌存储在会话中
	c.Set("csrf_token", token)
	return token, nil
}

// 从请求中获取CSRF令牌
func getCSRFToken(c *gin.Context, headerName string) (string, error) {
	// 首先从请求头中获取
	token := c.GetHeader(headerName)
	if token != "" {
		zap.L().Debug("从请求头中获取到CSRF令牌", zap.String("token", token))
		return token, nil
	}

	// 如果请求头中没有，尝试从表单中获取
	token = c.PostForm("_csrf")
	if token != "" {
		zap.L().Debug("从表单中获取到CSRF令牌", zap.String("token", token))
		return token, nil
	}

	// 如果表单中也没有，尝试从URL查询参数中获取
	token = c.Query("_csrf")
	if token != "" {
		zap.L().Debug("从URL查询参数中获取到CSRF令牌", zap.String("token", token))
		return token, nil
	}

	zap.L().Debug("未找到CSRF令牌")
	return "", errors.New("未找到CSRF令牌")
}

// 验证CSRF令牌
func validateCSRFToken(c *gin.Context, token string) (bool, error) {
	// 从会话中获取存储的令牌
	storedToken, exists := c.Get("csrf_token")
	if !exists {
		// 如果会话中没有，尝试从Cookie中获取
		var err error
		storedToken, err = c.Cookie("csrf_token")
		if err != nil {
			zap.L().Debug("会话和Cookie中都没有CSRF令牌")
			return false, errors.New("会话和Cookie中都没有CSRF令牌")
		}
	}

	// 比较令牌（使用安全的时间常量比较方法）
	storedTokenStr, ok := storedToken.(string)
	if !ok {
		zap.L().Error("存储的CSRF令牌类型错误")
		return false, errors.New("存储的CSRF令牌类型错误")
	}

	// 使用安全的时间常量比较方法
	valid := compareTokens(token, storedTokenStr)
	zap.L().Debug("验证CSRF令牌",
		zap.String("提供的令牌", token),
		zap.String("存储的令牌", storedTokenStr),
		zap.Bool("验证结果", valid))

	return valid, nil
}

// 使用安全的时间常量比较方法比较两个令牌
func compareTokens(a, b string) bool {
	if len(a) != len(b) {
		return false
	}

	// 使用异或操作和按位或操作来实现时间常量比较
	// 这种方法可以防止时间侧信道攻击
	var result byte
	for i := 0; i < len(a); i++ {
		result |= a[i] ^ b[i]
	}

	return result == 0
}

// 设置CSRF令牌Cookie
func SetCSRFCookie(c *gin.Context, token string, config CSRFConfig) {
	// 设置Cookie
	c.SetCookie(
		config.CookieName,
		token,
		config.CookieMaxAge,
		config.CookiePath,
		config.CookieDomain,
		config.CookieSecure,
		config.CookieHTTPOnly,
	)

	zap.L().Debug("设置CSRF Cookie",
		zap.String("name", config.CookieName),
		zap.String("token", token),
		zap.Int("maxAge", config.CookieMaxAge),
		zap.String("path", config.CookiePath),
		zap.String("domain", config.CookieDomain),
		zap.Bool("secure", config.CookieSecure),
		zap.Bool("httpOnly", config.CookieHTTPOnly))
}
